#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : microbiome_bracken_parser.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/9/30 16:51
@Description : 将Bracken的结果解析为机构化的数据
"""
import re
from argparse import ArgumentParser, Namespace
from pathlib import Path

import pandas as pd
import json


# =============================
# -------- 私有辅助函数 --------
# =============================
def _rank_sort_key(x):
    """
    按照标准对所有的涉及到的分类单元进行排序
    :param x:
    :return:
    """
    m = re.match(r"([A-Z]+)(\d*)", x)
    if m:
        letter, number = m.groups()
        return ["R", "K", "P", "C", "O", "F", "G", "S"].index(letter), (
            int(number) if number else 0
        )
    else:
        return 9999, 0  # 非标准 rank 放最后


def _assign_domain(info: str) -> str:
    """
    通过物种的分类信息判断物种是属于 细菌 古细菌 甄姬 病毒 寄生虫 还是其他
    :param info: 物种分类信息
    :return: 物种的域
    """
    if re.search(r"R\d*::Viruses", info):
        return "Viruses"
    if re.search(r"K\d*::Fungi", info):
        return "Fungi"
    if re.search(r"R\d*::Archaea", info):
        return "Archaea"
    if re.search(r"R\d*::Bacteria", info):
        return "Bacteria"
    # ********************************************************
    # -------- 注意！--------
    # 寄生虫并不是标准的分类单元，只能根据专家共识来判断哪些科属种是人体寄生虫
    # 难免会有遗漏，需要手动更新正则表达式
    # *********************************************************
    parasite_keywords = (
        r"protozoa|protista|apicomplexa|ciliophora|kinetoplastida|"
        r"plasmodium|theileria|babesia|toxoplasma|cryptosporidium|"
        r"entamoeba|trichomonas|giardia|balantidium|eimeria|leishmania|trypanosoma|naegleria|acytostelium|"
        r"animalia|nematoda|ascaris|enterobius|strongyloides|ancylostoma|"
        r"necator|trichuris|dracunculus|filarioidea|onchocerca|brugia|"
        r"platyhelminthes|trematoda|cestoda|schistosoma|fasciola|"
        r"clonorchis|opistorchis|diphyllobothrium|taenia|echinococcus|hymenolepis|"
        r"microsporidia|enterocytozoon|encephalitozoon"
    )
    if re.search(parasite_keywords, info.lower()):
        return "Parasite"

    # 不能分类的归为其他 包括一些动植物 藻类等
    return "Other"


def main(args: Namespace):
    """
    Bracken结果文件解析
    :param args: 命令行参数
    :return:
    """
    output_path = Path(args.output).absolute()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    bracken_path = Path(args.bracken).absolute()
    report_path = Path(args.report).absolute()

    # fastp的结果文件 用于获取总的下机reads数 计算RPM值
    fastp_path = Path(args.fastp).absolute()

    # 读取bracken文件
    report_df = pd.read_csv(
        report_path,
        sep=r"\t *",
        engine="python",
        header=None,
        names=["percent", "clade_reads", "taxon_reads", "rank", "tax_id", "name"],
    )

    # 去除物种名前后的空格
    report_df["name"] = report_df["name"].str.strip()

    # 将分类单元进行排序
    rank_order = report_df["rank"].unique().tolist()
    rank_order.sort(key=_rank_sort_key)

    # 对所有的分类单元进行初始化
    current_path = {r: "" for r in rank_order}

    # 输出记录
    records = []

    # 遍历数据框
    for _, row in report_df.iterrows():

        # 分类
        rank = row["rank"]

        # 物种名称
        name = row["name"]

        # 更新当前的rank
        current_path[rank] = name

        # 将本分级下的所有分级都清空
        current_rank = rank_order.index(rank)
        for r in rank_order[current_rank + 1 :]:
            current_path[r] = ""

        # 当遍历到物种的时将信息保存
        if rank.startswith("S"):
            info = ";;".join([f"{k}::{v}" for k, v in current_path.items() if v])
            # 获取当前物种的科
            family = current_path.get("F")
            # 获取当前物种的属
            genus = current_path.get("G")
            # 获取当前物种的种
            species = current_path.get("S")
            # 保存到记录中
            records.append(
                {
                    "name": name,
                    "family": family,
                    "genus": genus,
                    "species": species,
                    "info": info,
                }
            )

    # 转换为数据框
    info_df = pd.DataFrame(records)

    # 判断该物种是 细菌/古细菌/真菌/病毒/寄生虫/其他
    info_df["domain"] = info_df["info"].apply(_assign_domain)

    # 合并bracken结果信息
    bracken_df = pd.read_csv(bracken_path, sep="\t", encoding="utf-8")
    out_df = bracken_df.merge(info_df, on="name", how="left").fillna("")

    # 计算RPM值
    with fastp_path.open("r", encoding="utf-8") as f:
        fastp_dict = json.load(f)
    total_reads = int(fastp_dict["summary"]["before_filtering"]["total_reads"])
    out_df["RPM"] = out_df["new_est_reads"] / (total_reads / 1e6)
    out_df.to_csv(output_path, sep="\t", index=False)


if __name__ == "__main__":
    parser = ArgumentParser(description="将Bracken的结果文件解析为机构化文件")
    parser.add_argument(
        "-b", "--bracken", type=str, required=True, help="Bracken结果文件"
    )
    parser.add_argument(
        "-r", "--report", type=str, required=True, help="Bracken报告文件"
    )
    parser.add_argument(
        "-f", "--fastp", type=str, required=True, help="fastp报告文件"
    )
    parser.add_argument("-o", "--output", type=str, required=True, help="输出文件")
    main(parser.parse_args())